//
//  Implementation based on: http://www.eternallyconfuzzled.com/tuts/datastructures/jsw_tut_andersson.aspx
//

#import "AATree.h"

#import "DebugLogging.h"

@interface AATree() // private methods.

@property(retain) AATreeNode *root;
@property(assign) NSUInteger count;
@property(copy) NSComparator keyComparator;

/*!
 * @abstract                            Deletes the node bound to the specified key.
 * @discussion                          The node is deleted by looking up the node with the
 *                                                      specified key. The node is removed as in any binary
 *                                                      search tree, with the added functionallity that the
 *                                                      difference in levels between parent and child should
 *                                                      be at most one. If it is more, this is fixed, which may
 *                                                      lead to skew and split operations.
 */
- (AATreeNode *) __deleteNodeAtKey:(id)aKey atRoot:(AATreeNode *)aRoot;


/*!
 * @abstract                            Insert the specified data in the AA tree.
 * @discussion                          The data is inserted by looking up the correct leaf
 *                                                      node, setting the node as the left or right child of
 *                                                      the leaf. The function is recursive, so the skew and
 *                                                      split operations are performed on all parents of the
 *                                                      added node automatically. If a node with the same key
 *                                                      as the new node is found, the data of the node is
 *                                                      replaced with the new data.
 *
 * @param aRoot                         The root node to add the new node.
 * @param aNode                         The node to add.
 * @result                                      The possibly new root.
 */
- (AATreeNode *) __insertNode:(AATreeNode *)aNode atRoot:(AATreeNode *)aRoot;


/*!
 * @abstract                            Lock the tree for reading.
 * @discussion                          See the header file for more information on the current
 *                                                      implementation of thread safety.
 */
- (void) __lockForReading;


/*!
 * @abstract                            Lock the tree for writing.
 * @discussion                          See the header file for more information on the current
 *                                                      implementation of thread safety.
 */
- (void) __lockForWriting;


/*!
 * @abstract                            Retrieves the node bound to the specified key.
 * @discussion                          This function uses the key comparator, as specified on
 *                                                      initialization of the tree.
 *
 * @param aKey                          The key to look for.
 * @result                                      An AATreeNode pointer.
 */
- (AATreeNode *) __nodeAtKey:(id)aKey;


/*!
 * @abstract                            Retrieves the node which comes closest to the specified key.
 * @discussion                          This function uses the key comparator, as specified on
 *                                                      initialization of the tree. The returned node's key won't
 *                                                      surpass the specified key. If no node with a key lower than
 *                                                      the specified key can be found, nil is returned. This function
 *                                                      performs its operation recursive.
 *
 * @param aKey                          The key to look for.
 * @param aRoot                         The root to search from.
 * @result                                      An AATreeNode pointer.
 */
- (AATreeNode *) __nodeClosestToKey:(id)aKey atRoot:(AATreeNode *)aRoot;


/*!
 * @abstract                            Performs a recursive skew operation.
 * @discussion                          This function makes sure that every violation of the first
 *                                                      balance rule, which states that no left horizontal logical links
 *                                                      are allowed, are fixed. It does this by rotating right at the
 *                                                      parent of the left horizontal logical link.
 *
 * @param aRoot                         The root node to check.
 * @result                                      The possibly new root after the skew, or the same if nothing
 *                                                      has been skewed.
 */
- (AATreeNode *) __skew:(AATreeNode *)aRoot;


/*!
 * @abstract                            Performs a recursive split operation.
 * @discussion                          This function makes sure that every violation of the second
 *                                                      balance rule, which states that no two consecutive right
 *                                                      horizontal logical links are allowed, are fixed. It does this
 *                                                      by rotating left and increasing the level of the parent.
 *
 * @param aRoot                         The root node to check.
 * @result                                      The possibly new root after the split, or the same if nothing
 *                                                      has been split.
 */
- (AATreeNode *) __split:(AATreeNode *)aRoot;


/*!
 * @abstract                            Unlock the lock securing thread safety.
 */
- (void) __unlock;

@end


@implementation AATree

// --- --- --- --- --- --- --- --- --- --- --- --- --- --- --- --- --- --- --- --- ---
// -- public methods --
// --- --- --- --- --- --- --- --- --- --- --- --- --- --- --- --- --- --- --- --- ---

- (id) initWithKeyComparator:(NSComparator)aKeyComparator {

        if (self = [super init]) {
                keyComparator = [aKeyComparator copy];
                pthread_rwlock_init(&rwLock, NULL);
        }
        return self;
}


- (id) init {

        @throw [NSException exceptionWithName:@"MethodNotAllowedException" reason:@"Initialize an AATree with the 'initWithKeyComparator:' method." userInfo:nil];
}


- (id) copyWithZone:(NSZone *)zone {

        AATree *copy = [[AATree alloc] initWithKeyComparator:keyComparator];

        [self __lockForReading];
        copy.root = [[root copy] autorelease];
        copy.count = count;
        [self __unlock];

        return copy;
}

- (NSString *)description {
    return [NSString stringWithFormat:@"<%@: %p count=%@ root=%@>", self.class, self, @(self.count), self.root];
}

- (void)setCount:(NSUInteger)newCount {
    count = newCount;
}

- (NSUInteger) count {

        return count;
}



- (NSEnumerator *) keyEnumerator {

        NSMutableArray *keys = [NSMutableArray arrayWithCapacity:count];

        [self __lockForReading];
        [root addKeyToArray:keys];
        [self __unlock];

        return [keys objectEnumerator];
}


- (id) objectClosestToKey:(id)aKey {

        [self __lockForReading];
        id object = [self __nodeClosestToKey:aKey atRoot:root].data;
        [self __unlock];

        return object;
}


- (id) objectForKey:(id)aKey {

        [self __lockForReading];
        id data = [self __nodeAtKey:aKey].data;
        [self __unlock];

        return data;
}


- (void) print {

        [root printWithIndent:0];
}


- (void) removeObjectForKey:(id)aKey {

        [self __lockForWriting];
    changedNodes = [[NSMutableSet alloc] init];
    AATreeNode *prevRoot = self.root;
        self.root = [self __deleteNodeAtKey:aKey atRoot:root];
    if (prevRoot != self.root) {
        if (self.root) {
            [changedNodes addObject:self.root];
        }
        if (prevRoot) {
            [changedNodes addObject:prevRoot];
        }
    }
    if (changedNodes.count > 0) {
        [_delegate aaTree:self didChangeSubtreesAtNodes:changedNodes];
    }
    [changedNodes release];
    changedNodes = nil;
        [self __unlock];
}


- (void) setObject:(id)anObject forKey:(id)aKey {

        NSParameterAssert(anObject);
        NSParameterAssert(aKey);
        ITAssertWithMessage([aKey conformsToProtocol:@protocol(NSCopying)], @"The supplied key does not conform to the NSCopying protocol.");

        [self __lockForWriting];
    changedNodes = [[NSMutableSet alloc] init];
        AATreeNode *newNode = [[[AATreeNode alloc] initWithData:anObject boundToKey:[[aKey copy] autorelease]] autorelease];
    AATreeNode *prevRoot = self.root;
        self.root = [self __insertNode:newNode atRoot:root];
    if (prevRoot != self.root) {
        if (self.root) {
            [changedNodes addObject:self.root];
        }
        if (prevRoot) {
            [changedNodes addObject:prevRoot];
        }
    }
    if (changedNodes.count > 0) {
        [_delegate aaTree:self didChangeSubtreesAtNodes:changedNodes];
    }
    [changedNodes release];
    changedNodes = nil;
        [self __unlock];
}

- (void) notifyValueChangedForKey:(id)aKey {
    [_delegate aaTree:self didChangeValueAtNode:[self __nodeAtKey:aKey]];
}

- (NSArray *)pathFromNode:(AATreeNode *)node {
    NSMutableArray *parents = [NSMutableArray array];

    // Begin at the root of the tree.
    AATreeNode *current = root;

    // While still at a node, check whether we have found the correct node or
    // travel left or right.
    while (current) {
        [parents insertObject:current atIndex:0];
        NSComparisonResult compareResult = keyComparator(node.key, current.key);
        if (compareResult == NSOrderedSame)     return parents;
        else if (compareResult == NSOrderedAscending) current = current.left;
        else current = current.right;
    }

    // Nothing found, return nil.
    return nil;
}

- (void) dealloc
{
        [root release];
        [keyComparator release];
        [super dealloc];
}


// --- --- --- --- --- --- --- --- --- --- --- --- --- --- --- --- --- --- --- --- ---
// -- private methods --
// --- --- --- --- --- --- --- --- --- --- --- --- --- --- --- --- --- --- --- --- ---

@synthesize root;
@synthesize count;
@synthesize keyComparator;

- (AATreeNode *) __deleteNodeAtKey:(id)aKey atRoot:(AATreeNode *)aRoot {

        if (aRoot) {

                // If we found the correct node, remove it.
                NSComparisonResult compareResult = keyComparator(aKey, aRoot.key);
        if (compareResult == NSOrderedSame) {

            // Check whether we are at an easy to remove node (zero to one children) or
            // a more difficult node.
            if (aRoot.left && aRoot.right) {

                // Get the in-order predecessor (heir).
                AATreeNode *heir = aRoot.left;
                while (heir.right) heir = heir.right;

                // Replace the data.
                aRoot.key = heir.key;
                aRoot.data = heir.data;

                // Delete the in-order predecessor (heir).
                aRoot.left = [self __deleteNodeAtKey:aRoot.key atRoot:aRoot.left];
                [changedNodes addObject:aRoot];
                if (aRoot.left) {
                    [changedNodes addObject:aRoot.left];
                }
            } else {
                if (aRoot.left) {
                    aRoot.deleted = YES;
                    aRoot = aRoot.left;
                } else {
                    aRoot.deleted = YES;
                    aRoot = aRoot.right; // which could be nil.
                }
                assert(count > 0);
                count--;
            }

            // Otherwise, travel left or right.
                } else if (compareResult == NSOrderedAscending) {
            AATreeNode *prevLeft = aRoot.left;
                        aRoot.left = [self __deleteNodeAtKey:aKey atRoot:aRoot.left];
            if (prevLeft != aRoot.left) {
                [changedNodes addObject:aRoot];
                [changedNodes addObject:prevLeft];
            }
                } else {
            AATreeNode *prevRight = aRoot.right;
                        aRoot.right = [self __deleteNodeAtKey:aKey atRoot:aRoot.right];
            if (prevRight != aRoot.right) {
                [changedNodes addObject:aRoot];
                [changedNodes addObject:prevRight];
            }
                }

                // Check whether the levels or the children are not more than one
                // lower than the current.
                if (aRoot.left.level < aRoot.level - 1 || aRoot.right.level < aRoot.level - 1) {

                        // Decrease the level by one.
                        aRoot.level--;

                        // Decrease the right child's level also, when it is higher than its parent.
                        if (aRoot.right.level > aRoot.level) {
                aRoot.right.level = aRoot.level;
            }

            AATreeNode *prevRoot = aRoot;
                        aRoot = [self __skew:aRoot];
            if (aRoot != prevRoot) {
                [changedNodes addObject:prevRoot];
                [changedNodes addObject:aRoot];
            }
            prevRoot = aRoot;
                        aRoot = [self __split:aRoot];
            if (aRoot != prevRoot) {
                [changedNodes addObject:prevRoot];
                [changedNodes addObject:aRoot];
            }
                }
        }

        return aRoot;
}


- (AATreeNode *) __insertNode:(AATreeNode *)aNode atRoot:(AATreeNode *)aRoot {

        // If the root is not nil, we have not reached an empty child of a leaf node.
        if (aRoot) {

                // Decide which way to travel through the tree.
                NSComparisonResult compareResult = keyComparator(aNode.key, aRoot.key);

                // If the key of the new node is equal to the current root, just replace the data.
                if (compareResult == NSOrderedSame)     {
                        aRoot.data = aNode.data;
            [changedNodes addObject:aRoot];
            // Otherwise, travel left or right through the tree.
                } else {
                        if (compareResult == NSOrderedAscending) {
                aRoot.left = [self __insertNode:aNode atRoot:aRoot.left];
                [changedNodes addObject:aRoot];
                [changedNodes addObject:aRoot.left];
            } else {
                aRoot.right = [self __insertNode:aNode atRoot:aRoot.right];
                [changedNodes addObject:aRoot];
                [changedNodes addObject:aRoot.right];
            }

                        // After the node has been added, skew and split the (possibly new) root.
                        // Because of the recursive nature of this function, all parents of the
                        // new node will get skewed and split, all the way up to the root of the tree.
            AATreeNode *prevRoot = aRoot;
                        aRoot = [self __skew:aRoot];
            if (aRoot != prevRoot) {
                [changedNodes addObject:aRoot];
                [changedNodes addObject:prevRoot];
            }
            prevRoot = aRoot;
                        aRoot = [self __split:aRoot];
            if (aRoot != prevRoot) {
                [changedNodes addObject:aRoot];
                [changedNodes addObject:prevRoot];
            }
                }

        // Otherwise, insert the node.
        } else {
                aRoot = aNode;
                count++;
        }

        return aRoot;
}


- (void) __lockForReading {

        pthread_rwlock_rdlock(&rwLock);
}


- (void) __lockForWriting {

        pthread_rwlock_wrlock(&rwLock);
}


- (AATreeNode *) __nodeAtKey:(id)aKey {

        // Begin at the root of the tree.
        AATreeNode *current = root;

        // While still at a node, check whether we have found the correct node or
        // travel left or right.
        while (current) {
                NSComparisonResult compareResult = keyComparator(aKey, current.key);
                if (compareResult == NSOrderedSame)     return current;
                else if (compareResult == NSOrderedAscending) current = current.left;
                else current = current.right;
        }

        // Nothing found, return nil.
        return nil;
}


- (AATreeNode *) __nodeClosestToKey:(id)aKey atRoot:(AATreeNode *)aRoot {

        // Start with no result.
        AATreeNode *result = nil;

        // If we are still at a node, compare it to the specified key.
        if (aRoot) {
                NSComparisonResult compareResult = keyComparator(aKey, aRoot.key);

                // If the keys are equal, we have found an exact match and we are done.
                if (compareResult == NSOrderedSame)     result = aRoot;

                // Otherwise, travel left or right until a leaf node is surpassed.
                else if (compareResult == NSOrderedAscending) result = [self __nodeClosestToKey:aKey atRoot:aRoot.left];
                else result = [self __nodeClosestToKey:aKey atRoot:aRoot.right];

                // If no result has been found lower in the tree, test whether this node
                // is the closest.
                if (!result && compareResult == NSOrderedDescending) result = aRoot;
        }

        return result;
}


- (AATreeNode *) __skew:(AATreeNode *)aRoot {

        if (aRoot) {

                // Check for a logical horizontal left link.
                if (aRoot.left.level == aRoot.level) {
                        // Perform a right rotation.
                        AATreeNode *save = aRoot;
                        aRoot = aRoot.left;
                        save.left = aRoot.right;
                        aRoot.right = save;
                }
                // Skew the right side of the (new) root.
                aRoot.right = [self __skew:aRoot.right];
        [changedNodes addObject:aRoot];
        }

        return aRoot;
}


- (AATreeNode *) __split:(AATreeNode *)aRoot {

        // Check for a consecutive logical horizontal right link.
        if (aRoot && aRoot.right.right.level == aRoot.level) {
                // Perform a left rotation.
                AATreeNode *save = aRoot;
                aRoot = aRoot.right;
                save.right = aRoot.left;
                aRoot.left = save;

                // Increase the level of the new root.
                aRoot.level++;

                // Split the right side of the new root.
                aRoot.right = [self __split:aRoot.right];
        [changedNodes addObject:aRoot];
        }

        return aRoot;
}


- (void) __unlock {

        pthread_rwlock_unlock(&rwLock);
}

@end
